from langchain_openai import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

llm = ChatOpenAI(
    api_key="sk-0e687ddcf0164a6fb66c1096447223c4",  # 阿里百炼大模型API获取：https://bailian.console.aliyun.com/?apiKey=1#/api-key
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", # 使用文档：https://help.aliyun.com/zh/model-studio/getting-started/what-is-model-studio
    model="qwen-plus",
    temperature=0.7
    # other params...
)
# 初始提示
initial_instructions = "你是一个友好的聊天机器人，尽可能简洁地回答问题。"

# 对话记忆
memory = ConversationBufferMemory()

# 用于调用模型的提示模板
call_prompt_template = PromptTemplate.from_template(
    "{instructions}\n历史对话：{history}\n用户输入：{input}"
)

# 用于更新指令的提示模板
update_prompt_template = PromptTemplate.from_template(
    "当前指令：{instructions}\n最近对话：{conversation}\n请根据以上信息更新指令："
)

# 调用模型的函数
def call_model(instructions, user_input):
    memory_vars = memory.load_memory_variables({})
    chain = call_prompt_template | llm
    response = chain.invoke({
        "instructions": instructions,
        "history": memory_vars.get("history", ""),
        "input": user_input
    })
    return response.content

# 更新指令的函数
def update_instructions(instructions):
    conversation = memory.buffer
    chain = update_prompt_template | llm
    response = chain.invoke({
        "instructions": instructions,
        "conversation": conversation
    })
    return response.content

# 主对话循环
def main():
    instructions = initial_instructions
    print("欢迎使用AI助手！输入'quit'退出对话。")
    
    while True:
        # 获取用户输入
        user_input = input("> ")
        if user_input.lower() == 'quit':
            print("感谢使用，再见！")
            break
            
        # 调用模型获取回复
        response = call_model(instructions, user_input)
        print(f"AI: {response}")
        
        # 询问用户反馈
        print("\n您对这个回答满意吗？如果需要调整，请告诉我。(直接回车继续对话)")
        feedback = input("反馈> ")
        
        # 如果用户提供了反馈，更新指令
        if feedback.strip():
            memory.chat_memory.add_user_message(feedback)
            instructions = update_instructions(instructions)
            print("已根据您的反馈更新对话策略。")
            print(f"当前指令：{instructions}")

if __name__ == "__main__":
    main()